{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/dlachinov/pytorch_dl_pc/lib/python3.6/site-packages/dicom/__init__.py:53: UserWarning: \n",
      "This code is using an older version of pydicom, which is no longer \n",
      "maintained as of Jan 2017.  You can access the new pydicom features and API \n",
      "by installing `pydicom` from PyPI.\n",
      "See 'Transitioning to pydicom 1.x' section at pydicom.readthedocs.org \n",
      "for more information.\n",
      "\n",
      "  warnings.warn(msg)\n"
     ]
    }
   ],
   "source": [
    "import loader_helper\n",
    "import os\n",
    "import numpy as np\n",
    "import tqdm\n",
    "import nibabel as nii\n",
    "from scipy.ndimage import zoom"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = '../../data/MICCAI_BraTS_2019_Data_Training/ALL'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "335\n"
     ]
    }
   ],
   "source": [
    "files = [f for f in os.listdir(data_path) if os.path.isdir(os.path.join(data_path,f))]\n",
    "files.sort()\n",
    "\n",
    "print(len(files))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "data, annotation, affine = loader_helper.read_multimodal(data_path, files[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_bbox(data):\n",
    "    bboxes = np.stack([loader_helper.bbox3(d) for d in data],axis=0)\n",
    "    return np.stack([np.min(bboxes[:,0],axis=0),np.max(bboxes[:,1],axis=0)],axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 335/335 [01:22<00:00,  4.11it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(335, 2, 3)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "bboxes = []\n",
    "affines = []\n",
    "for f in tqdm.tqdm(files):\n",
    "    data, annotation, affine = loader_helper.read_multimodal(data_path, f, False)\n",
    "    affines.append(affine)\n",
    "    bboxes.append(get_bbox(data))\n",
    "bboxes = np.stack(bboxes,axis=0)\n",
    "print(bboxes.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[156 188 148]\n"
     ]
    }
   ],
   "source": [
    "sizes = [b[1] - b[0] for b in bboxes]\n",
    "print(np.max(sizes,axis=0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[160, 192, 160]\n"
     ]
    }
   ],
   "source": [
    "crop_size = [loader_helper.closest_to_k(i,16) for i in np.max(sizes,axis=0)]\n",
    "print(crop_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 335/335 [05:13<00:00,  1.02it/s]\n"
     ]
    }
   ],
   "source": [
    "min = np.array([0,0,0])\n",
    "max = np.array([240, 240, 155])\n",
    "\n",
    "output_dir = '../../data/2019_cropped'\n",
    "\n",
    "crop_size = np.array(crop_size)\n",
    "\n",
    "#resample_to_size = (64,64,64)\n",
    "#scale_factor = np.array(crop_size) / np.array(resample_to_size)\n",
    "\n",
    "for idx, f in enumerate(tqdm.tqdm(files)):\n",
    "    data, annotation, affine = loader_helper.read_multimodal(data_path, f, True)\n",
    "    b = get_bbox(data)\n",
    "    size = b[1] - b[0]\n",
    "    \n",
    "    output = np.zeros(shape=(data.shape[0],)+tuple(crop_size))\n",
    "    out_annotation = np.zeros(shape=tuple(crop_size))\n",
    "    \n",
    "    diff = np.array(crop_size) - np.array(size)\n",
    "    low = diff // 2\n",
    "    high = low - diff\n",
    "    bbox = b - np.stack([low,high])\n",
    "\n",
    "    index_input_min = np.maximum(bbox[0],min)\n",
    "    index_input_max = np.minimum(bbox[1],max)\n",
    "    \n",
    "    size = index_input_max - index_input_min\n",
    "    \n",
    "\n",
    "    index_output_min = crop_size//2 - size//2 \n",
    "    index_output_max = crop_size//2 + size - size//2\n",
    "\n",
    "    \n",
    "    output[:,index_output_min[0]:index_output_max[0],index_output_min[1]:index_output_max[1],index_output_min[2]:index_output_max[2]] =\\\n",
    "    data[:,index_input_min[0]:index_input_max[0],index_input_min[1]:index_input_max[1],index_input_min[2]:index_input_max[2]]\n",
    "    \n",
    "    out_annotation[index_output_min[0]:index_output_max[0],index_output_min[1]:index_output_max[1],index_output_min[2]:index_output_max[2]] =\\\n",
    "    annotation[index_input_min[0]:index_input_max[0],index_input_min[1]:index_input_max[1],index_input_min[2]:index_input_max[2]]\n",
    "    \n",
    "    \n",
    "\n",
    "    suffixes = ['_t1.nii.gz','_t1ce.nii.gz','_t2.nii.gz','_flair.nii.gz']\n",
    "    \n",
    "    os.makedirs(name=os.path.join(output_dir,f), exist_ok=True)\n",
    "    \n",
    "    #affine = np.diag([affine[i,i]*scale_factor[i] for i in range(3)]+[1])    \n",
    "    \n",
    "    for jdx, d in enumerate(output):\n",
    "        #d = zoom(d, 1/scale_factor, order=3, mode='constant', cval=0)\n",
    "    \n",
    "        output_header = nii.Nifti1Image(d.astype(np.float32), affine)\n",
    "        nii.save(output_header, os.path.join(output_dir,f,f+suffixes[jdx]))\n",
    "        \n",
    "    \n",
    "    #out_annotation = zoom(out_annotation, 1/scale_factor, order=0, mode='constant', cval=0)\n",
    "    out_annotation[out_annotation==4] = 3\n",
    "    output_header = nii.Nifti1Image(out_annotation.astype(np.uint8), affine)\n",
    "    nii.save(output_header, os.path.join(output_dir,f,f+'_seg.nii.gz'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
